import pandas as pd
import numpy as np
from scipy.fft import fft, fftfreq
from scipy.signal import find_peaks


def data_process(df):
    # 删除空值行
    df.dropna(inplace=True)

    # # 删除异常值
    # df = df[(50000 <= df['频率，Hz']) & (df['频率，Hz'] <= 500000)]

    # 使用map()方法将列中的值映射到新值
    # 1表示正弦波，2表示三角波，3表示梯形波
    mapping_dict = {
        '正弦波': 1,
        '三角波': 2,
        '梯形波': 3
    }

    df['励磁波形'] = df['励磁波形'].map(mapping_dict)
    return df


def feature_extract(df):
    # 提取前四列
    meta_data = df.iloc[:, :4]
    # 提取磁通密度数据
    data = df.iloc[:, 4:]
    # 初始化一个空的列表来存储特征
    features = []

    # 遍历每一行数据，计算特征
    for index, row in data.iterrows():
        signal = row.values

        # 计算时域特征
        mean_value = np.mean(signal)  # 均值
        variance = np.var(signal)  # 方差
        max_value = np.max(signal)  # 最大值
        min_value = np.min(signal)  # 最小值
        peak_to_peak = max_value - min_value  # 峰峰值
        kurtosis = pd.Series(signal).kurtosis()  # 峰度
        skew = pd.Series(signal).skew()  # 偏度

        # 计算频域特征
        N = len(signal)
        sampling_rate = 1024  # 采样频率
        # 傅里叶变换
        fft_result = fft(signal)
        # 频率轴
        frequencies = fftfreq(N, 1 / sampling_rate)
        # 频谱
        fft_magnitude = np.abs(fft_result)
        # 频谱质心
        spectral_centroid = np.sum(frequencies * fft_magnitude) / np.sum(fft_magnitude)
        # 频谱熵
        power_spectrum = (fft_magnitude ** 2) / np.sum(fft_magnitude ** 2)  # 归一化
        spectral_entropy = -np.sum(power_spectrum * np.log2(power_spectrum + 1e-12))
        # 带宽
        bandwidth = np.sqrt(np.sum(((frequencies - spectral_centroid) ** 2) * fft_magnitude) / np.sum(fft_magnitude))
        # 频谱峰值数
        peaks, _ = find_peaks(fft_magnitude)
        num_peaks = len(peaks)
        # 频率峰值间距
        if num_peaks > 1:
            peak_distances = np.diff(frequencies[peaks])
            avg_peak_distance = np.mean(peak_distances)
        else:
            avg_peak_distance = 0

        # 将特征添加到列表中
        features.append({
            '均值': mean_value,
            '方差': variance,
            '最大值': max_value,
            '最小值': min_value,
            '峰峰值': peak_to_peak,
            '峰度': kurtosis,
            '偏度': skew,
            '频谱质心': spectral_centroid,
            '频谱熵': spectral_entropy,
            '带宽': bandwidth,
            '频谱峰值数': num_peaks,
            '频率峰值间距': avg_peak_distance,
        })

    # 将特征列表转换为DataFrame
    features_df = pd.DataFrame(features)
    extracted_df = pd.concat([meta_data, features_df], axis=1)

    return extracted_df


if __name__ == '__main__':
    material_list = ['材料1', '材料2', '材料3', '材料4']
    file_path = '../data/Q1_data.xlsx'

    # 初始化一个excel表格
    pd.DataFrame().to_excel(file_path, sheet_name=material_list[0])

    for i in range(len(material_list)):
        df = pd.read_excel('../2024年中国研究生数学建模竞赛赛题/C题/附件一（训练集）.xlsx', sheet_name=material_list[i])
        print(f'正在提取{material_list[i]}特征')
        mapped_data = data_process(df)
        extracted_data = feature_extract(mapped_data)
        with pd.ExcelWriter(file_path, engine='openpyxl', mode='a', if_sheet_exists='replace') as writer:
            extracted_data.to_excel(writer, sheet_name=material_list[i], index=False)
        print(f'{material_list[i]}特征提取完成')

